import tensorflow as tf
import numpy as np
import keras
from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D
from keras.models import Sequential
from keras.datasets import mnist
import matplotlib.pyplot as plt

net = keras.models.Sequential([
    keras.layers.Conv2D(filters=96, kernel_size=(11, 11), strides=(4, 4), padding='valid', activation='relu'),
    keras.layers.MaxPool2D(pool_size=(3, 3), strides=(2, 2), padding='valid'),
    keras.layers.Conv2D(filters=256, kernel_size=(5, 5), strides=(1, 1), padding='same', activation='relu'),
    keras.layers.MaxPool2D(pool_size=(3, 3), strides=(2, 2), padding='valid'),
    keras.layers.Conv2D(filters=384, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'),
    keras.layers.Conv2D(filters=384, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'),
    keras.layers.Conv2D(filters=256, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'),
    keras.layers.MaxPool2D(pool_size=(3, 3), strides=(2, 2), padding='valid'),
    keras.layers.Flatten(),
    keras.layers.Dense(4096, activation='relu'),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(4096, activation='relu'),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(1000, activation='softmax')
])
X = tf.random.uniform((1, 227, 227, 1))
y = net(X)
print(net.summary())
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
# 维度调整
train_images = np.reshape(train_images,(train_images.shape[0],train_images.shape[1],train_images.shape[2],1))
test_images = np.reshape(test_images,(test_images.shape[0],test_images.shape[1],test_images.shape[2],1))


def get_train(size):
    index = np.random.randint(0, 60000, size)
    resized_images = tf.image.resize(train_images[index], (227, 227))
    return resized_images.numpy(), train_labels[index]


def get_test(size):
    index = np.random.randint(0, 10000, size)
    resized_images = tf.image.resize(test_images[index], (227, 227))
    return resized_images.numpy(), test_labels[index]


train_images, train_labels = get_train(256)
test_images, test_labels = get_test(256)

plt.imshow(train_images[0].astype(np.uint8).squeeze(), cmap="gray")
plt.show()
net.compile(optimizer=keras.optimizers.SGD(learning_rate=0.01), loss=keras.losses.SparseCategoricalCrossentropy(),
            metrics=['accuracy'])
net.fit(train_images, train_labels, epochs=10, batch_size=128, validation_split=0.1, verbose=1)
net.evaluate(test_images, test_labels, verbose=1)
